import cv2
import streamlit as st
from PIL import Image

from utils.Detect import *
from utils.LoadModel import *

st.logo(Config.logo)

mission = '安全帽检测'
st.header(mission)
print('-' * 20, mission, '-' * 20)
person_model = load_person_model()
hat_model = load_hat_model()

func = st.sidebar.selectbox(
    '选择数据源',
    ['图片检测', '摄像头']
)

# 添加启动和停止按钮
if 'running' not in st.session_state:
    st.session_state['running'] = False  # 初始化运行状态为False

if func == '摄像头':
    if st.sidebar.button('开启摄像头'):
        st.session_state['running'] = True

    if st.sidebar.button('关闭摄像头'):
        st.session_state['running'] = False

# 图片与模型
if func == '图片检测':
    uploaded_img = st.file_uploader("上传照片进行检测", type=["jpg", "jpeg", "png"])
    col1, col2 = st.columns(2)

    if uploaded_img is not None:
        img = Image.open(uploaded_img)
        with col1:
            st.image(img)

        results = detect(person_model, img)
        img_draw = draw_img(person_model, results, img)

        results = detect(hat_model, img)
        img_draw = draw_img(hat_model, results, img_draw)

        path = Config.hisPic + f'{time.time()}' + '.jpg'
        img_draw.save(path)

        save_data(mission, path)

        with col2:
            st.image(img_draw)

elif func == '摄像头':
    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        st.error("无法打开摄像头，请检查设备连接。")
    else:
        video_placeholder = st.empty()
        while st.session_state.running and cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                st.error("无法获取帧。")
                break

            results = detect(person_model, frame)
            frame_drawed = draw_img(person_model, results, Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)))

            results = detect(hat_model, frame)
            frame_drawed = draw_img(hat_model, results, frame_drawed)
            
            video_placeholder.image(frame_drawed, channels="RGB", use_container_width=True)

        cap.release()
